// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
// or more contributor license agreements. Licensed under the Elastic License 2.0;
// you may not use this file except in compliance with the Elastic License 2.0.

package sampling

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/gofrs/uuid/v5"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/otel/metric/noop"

	"github.com/elastic/apm-data/model/modelpb"
)

func TestTraceGroupsPolicies(t *testing.T) {
	makeTransaction := func(serviceName, serviceEnvironment, traceOutcome, traceName string) *modelpb.APMEvent {
		return &modelpb.APMEvent{
			Service: &modelpb.Service{
				Name:        serviceName,
				Environment: serviceEnvironment,
			},
			Event: &modelpb.Event{
				Outcome: traceOutcome,
			},
			Trace: &modelpb.Trace{Id: uuid.Must(uuid.NewV4()).String()},
			Transaction: &modelpb.Transaction{
				Type: "type",
				Name: traceName,
				Id:   uuid.Must(uuid.NewV4()).String(),
			},
		}
	}
	makePolicy := func(sampleRate float64, serviceName, serviceEnvironment, traceOutcome, traceName string) Policy {
		return Policy{
			SampleRate: sampleRate,
			PolicyCriteria: PolicyCriteria{
				ServiceName:        serviceName,
				ServiceEnvironment: serviceEnvironment,
				TraceName:          traceName,
				TraceOutcome:       traceOutcome,
			},
		}
	}

	const (
		staticServiceName  = "static-service"
		dynamicServiceName = "dynamic-service"
	)
	policies := []Policy{
		makePolicy(0.4, staticServiceName, "production", "error", "GET /healthcheck"),
		makePolicy(0.3, staticServiceName, "production", "error", ""),
		makePolicy(0.2, staticServiceName, "production", "", ""),
		makePolicy(0.1, staticServiceName, "", "", ""),
	}

	// Clone the policies without ServiceName set, so we have identical catch-all policies
	// that will match the dynamic service name.
	for _, policy := range policies[:] {
		policy.ServiceName = ""
		policies = append(policies, policy)
	}
	groups := newTraceGroups(noop.Meter{}, policies, 1000, 1.0)

	assertSampleRate := func(sampleRate float64, serviceName, serviceEnvironment, traceOutcome, traceName string) {
		tx := makeTransaction(serviceName, serviceEnvironment, traceOutcome, traceName)
		const N = 1000
		for i := 0; i < N; i++ {
			if _, err := groups.sampleTrace(tx); err != nil {
				t.Fatal(err)
			}
		}
		sampled := groups.finalizeSampledTraces(nil)
		assert.Len(t, sampled, int(sampleRate*N))
	}

	for _, serviceName := range []string{staticServiceName, dynamicServiceName} {
		assertSampleRate(0.1, serviceName, "testing", "error", "GET /healthcheck")
		assertSampleRate(0.2, serviceName, "production", "success", "GET /healthcheck")
		assertSampleRate(0.3, serviceName, "production", "error", "GET /")
		assertSampleRate(0.4, serviceName, "production", "error", "GET /healthcheck")
	}
}

func TestTraceGroupsMax(t *testing.T) {
	const (
		maxDynamicServices    = 100
		ingestRateCoefficient = 1.0
	)
	policies := []Policy{{SampleRate: 1.0}}
	groups := newTraceGroups(noop.Meter{}, policies, maxDynamicServices, ingestRateCoefficient)

	for i := 0; i < maxDynamicServices; i++ {
		serviceName := fmt.Sprintf("service_group_%d", i)
		for i := 0; i < minReservoirSize; i++ {
			admitted, err := groups.sampleTrace(&modelpb.APMEvent{
				Service: &modelpb.Service{
					Name: serviceName,
				},
				Trace: &modelpb.Trace{Id: uuid.Must(uuid.NewV4()).String()},
				Transaction: &modelpb.Transaction{
					Type: "type",
					Name: "whatever",
					Id:   uuid.Must(uuid.NewV4()).String(),
				},
			})
			require.NoError(t, err)
			assert.True(t, admitted)
		}
	}

	admitted, err := groups.sampleTrace(&modelpb.APMEvent{
		Trace: &modelpb.Trace{Id: uuid.Must(uuid.NewV4()).String()},
		Transaction: &modelpb.Transaction{
			Type: "type",
			Name: "overflow",
			Id:   uuid.Must(uuid.NewV4()).String(),
		},
	})
	assert.Equal(t, errTooManyTraceGroups, err)
	assert.False(t, admitted)
}

func TestTraceGroupReservoirResize(t *testing.T) {
	const (
		maxDynamicServices    = 1
		ingestRateCoefficient = 0.75
	)
	policies := []Policy{{SampleRate: 0.2}}
	groups := newTraceGroups(noop.Meter{}, policies, maxDynamicServices, ingestRateCoefficient)

	sendTransactions := func(n int) {
		for i := 0; i < n; i++ {
			groups.sampleTrace(&modelpb.APMEvent{
				Trace: &modelpb.Trace{Id: "0102030405060708090a0b0c0d0e0f10"},
				Transaction: &modelpb.Transaction{
					Type: "type",
					Id:   "0102030405060708",
				},
			})
		}
	}

	// All groups start out with a reservoir size of 1000.
	sendTransactions(10000)
	assert.Len(t, groups.finalizeSampledTraces(nil), 1000) // initial reservoir size

	// We sent 10000 initially, and we send 20000 each subsequent iteration.
	// The number of sampled trace IDs will converge on 4000 (0.2*20000).
	for i, expected := range []int{
		2000, // 0.2 * 10000 (initial ingest rate)
		3500, // 0.2 * (0.25*10000 + 0.75*20000)
		3875, // 0.2 * (0.25*17500 + 0.75*20000)
		3969, // etc.
		3993,
		3999,
		4000,
		4000,
	} {
		sendTransactions(20000)
		traces := groups.finalizeSampledTraces(nil)
		assert.Len(t, traces, expected, "iteration %d expected len %d actual len %d", i, expected, len(traces))
	}
}

func TestTraceGroupReservoirResizeMinimum(t *testing.T) {
	const (
		maxDynamicServices    = 1
		ingestRateCoefficient = 1.0
	)
	policies := []Policy{{SampleRate: 0.1}}
	groups := newTraceGroups(noop.Meter{}, policies, maxDynamicServices, ingestRateCoefficient)

	sendTransactions := func(n int) {
		for i := 0; i < n; i++ {
			groups.sampleTrace(&modelpb.APMEvent{
				Trace: &modelpb.Trace{Id: "0102030405060708090a0b0c0d0e0f10"},
				Transaction: &modelpb.Transaction{
					Type: "type",
					Id:   "0102030405060708",
				},
			})
		}
	}

	sendTransactions(10000)
	assert.Len(t, groups.finalizeSampledTraces(nil), 1000) // initial reservoir size

	// The reservoir would normally be resized to fit the desired sampling
	// rate, but will never be resized to less than the minimum (1000).
	sendTransactions(1000)
	assert.Len(t, groups.finalizeSampledTraces(nil), 100)

	sendTransactions(10000)
	assert.Len(t, groups.finalizeSampledTraces(nil), 1000) // min reservoir size
}

func TestTraceGroupsRemoval(t *testing.T) {
	const (
		maxDynamicServices    = 2
		ingestRateCoefficient = 1.0
	)
	policies := []Policy{
		{PolicyCriteria: PolicyCriteria{ServiceName: "defined"}, SampleRate: 0.5},
		{SampleRate: 0.5},
		{PolicyCriteria: PolicyCriteria{ServiceName: "defined_later"}, SampleRate: 0.5},
	}
	groups := newTraceGroups(noop.Meter{}, policies, maxDynamicServices, ingestRateCoefficient)

	for i := 0; i < 10000; i++ {
		_, err := groups.sampleTrace(&modelpb.APMEvent{
			Service:     &modelpb.Service{Name: "many"},
			Transaction: &modelpb.Transaction{Type: "type"},
		})
		assert.NoError(t, err)
	}
	_, err := groups.sampleTrace(&modelpb.APMEvent{
		Service:     &modelpb.Service{Name: "few"},
		Transaction: &modelpb.Transaction{Type: "type"},
	})
	assert.NoError(t, err)

	_, err = groups.sampleTrace(&modelpb.APMEvent{
		Service:     &modelpb.Service{Name: "another"},
		Transaction: &modelpb.Transaction{Type: "type"},
	})
	assert.Equal(t, errTooManyTraceGroups, err)

	// When there is a policy with an explicitly defined service name, that
	// will not be affected by the limit...
	_, err = groups.sampleTrace(&modelpb.APMEvent{
		Service:     &modelpb.Service{Name: "defined"},
		Transaction: &modelpb.Transaction{Type: "type"},
	})
	assert.NoError(t, err)

	// ...unless the policy with an explicitly defined service name comes after
	// a matching dynamic policy.
	_, err = groups.sampleTrace(&modelpb.APMEvent{
		Service:     &modelpb.Service{Name: "defined_later"},
		Transaction: &modelpb.Transaction{Type: "type"},
	})
	assert.Equal(t, errTooManyTraceGroups, err)

	// Finalizing should remove the "few" trace group, since its reservoir
	// size is at the minimum, and the number of groups is at the maximum.
	groups.finalizeSampledTraces(nil)

	// We should now be able to add another trace group.
	_, err = groups.sampleTrace(&modelpb.APMEvent{
		Service:     &modelpb.Service{Name: "another"},
		Transaction: &modelpb.Transaction{Type: "type"},
	})
	assert.NoError(t, err)
}

func TestTraceGroupsRemovalConcurrent(t *testing.T) {
	// Ensure that trace groups removal does not race with sampleTrace
	const (
		maxDynamicServices    = 2
		ingestRateCoefficient = 1.0
	)
	policies := []Policy{
		{SampleRate: 1},
	}
	groups := newTraceGroups(noop.Meter{}, policies, maxDynamicServices, ingestRateCoefficient)

	wg := sync.WaitGroup{}
	wg.Add(2)
	// create a dynamic group first
	_, err := groups.sampleTrace(&modelpb.APMEvent{
		Service:     &modelpb.Service{Name: "many"},
		Transaction: &modelpb.Transaction{Type: "type"},
	})
	assert.NoError(t, err)
	go func() {
		for i := 0; i < 10; i++ {
			_, err := groups.sampleTrace(&modelpb.APMEvent{
				Service:     &modelpb.Service{Name: "many"},
				Transaction: &modelpb.Transaction{Type: "type"},
			})
			assert.NoError(t, err)
		}
		wg.Done()
	}()
	go func() {
		groups.finalizeSampledTraces(nil)
		wg.Done()
	}()
	wg.Wait()
}

func BenchmarkTraceGroups(b *testing.B) {
	const (
		maxDynamicServices    = 1000
		ingestRateCoefficient = 1.0
	)
	policies := []Policy{{SampleRate: 1.0}}
	groups := newTraceGroups(noop.Meter{}, policies, maxDynamicServices, ingestRateCoefficient)

	b.RunParallel(func(pb *testing.PB) {
		// Transaction identifiers are different for each goroutine, simulating
		// multiple agentss. This should demonstrate low contention.
		//
		// Duration is non-zero to ensure transactions have a non-zero chance of
		// being sampled.
		tx := modelpb.APMEvent{
			Event: &modelpb.Event{Duration: uint64(time.Second)},
			Transaction: &modelpb.Transaction{
				Type: "type",
				Name: uuid.Must(uuid.NewV4()).String(),
			},
		}
		for pb.Next() {
			groups.sampleTrace(&tx)
			tx.Event.Duration = tx.Event.Duration + uint64(time.Second)
		}
	})
}
